"""
Database session management for GCR Solver Manager.
"""

import logging
from contextlib import asynccontextmanager, contextmanager
from typing import AsyncGenerator, Generator, Optional
from sqlalchemy import create_engine, event
from sqlalchemy.ext.asyncio import (
    AsyncSession, 
    async_sessionmaker, 
    create_async_engine,
    AsyncEngine
)
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import StaticPool

from ..core.config import get_config
from ..core.exceptions import DatabaseError
from .models import Base

logger = logging.getLogger(__name__)


class DatabaseManager:
    """Database connection and session manager."""
    
    def __init__(self, database_url: Optional[str] = None):
        """
        Initialize database manager.
        
        Args:
            database_url: Database URL. If None, uses config.
        """
        self.config = get_config()
        self.database_url = database_url or self.config.database.url
        
        # Determine if we're using async or sync engine
        self.is_async = self.database_url.startswith(('postgresql+asyncpg://', 'sqlite+aiosqlite://'))
        
        self._engine = None
        self._async_engine = None
        self._session_maker = None
        self._async_session_maker = None
        
        self._initialize_engine()
    
    def _initialize_engine(self):
        """Initialize database engine and session makers."""
        try:
            if self.is_async:
                self._initialize_async_engine()
            else:
                self._initialize_sync_engine()
        except Exception as e:
            raise DatabaseError(f"Failed to initialize database engine: {e}")
    
    def _initialize_sync_engine(self):
        """Initialize synchronous engine."""
        engine_kwargs = {
            'echo': self.config.database.echo,
            'pool_size': self.config.database.pool_size,
            'max_overflow': self.config.database.max_overflow,
            'pool_timeout': self.config.database.pool_timeout,
            'pool_recycle': self.config.database.pool_recycle,
        }
        
        # Special handling for SQLite
        if self.database_url.startswith('sqlite'):
            engine_kwargs.update({
                'poolclass': StaticPool,
                'connect_args': {'check_same_thread': False},
                'pool_pre_ping': True,
            })
            # Remove pool settings that don't apply to SQLite
            for key in ['pool_size', 'max_overflow', 'pool_timeout', 'pool_recycle']:
                engine_kwargs.pop(key, None)
        
        self._engine = create_engine(self.database_url, **engine_kwargs)
        
        # Add connection event listeners
        @event.listens_for(self._engine, "connect")
        def set_sqlite_pragma(dbapi_connection, connection_record):
            """Set SQLite pragmas for better performance and reliability."""
            if self.database_url.startswith('sqlite'):
                cursor = dbapi_connection.cursor()
                cursor.execute("PRAGMA foreign_keys=ON")
                cursor.execute("PRAGMA journal_mode=WAL")
                cursor.execute("PRAGMA synchronous=NORMAL")
                cursor.execute("PRAGMA temp_store=MEMORY")
                cursor.execute("PRAGMA mmap_size=268435456")  # 256MB
                cursor.close()
        
        self._session_maker = sessionmaker(
            bind=self._engine,
            autocommit=False,
            autoflush=False,
            expire_on_commit=False
        )
        
        logger.info(f"Initialized synchronous database engine: {self.database_url}")
    
    def _initialize_async_engine(self):
        """Initialize asynchronous engine."""
        # Convert sync URL to async URL if needed
        async_url = self.database_url
        if async_url.startswith('postgresql://'):
            async_url = async_url.replace('postgresql://', 'postgresql+asyncpg://')
        elif async_url.startswith('sqlite:///'):
            async_url = async_url.replace('sqlite:///', 'sqlite+aiosqlite:///')
        
        engine_kwargs = {
            'echo': self.config.database.echo,
            'pool_size': self.config.database.pool_size,
            'max_overflow': self.config.database.max_overflow,
            'pool_timeout': self.config.database.pool_timeout,
            'pool_recycle': self.config.database.pool_recycle,
        }
        
        # Special handling for SQLite
        if async_url.startswith('sqlite+aiosqlite'):
            engine_kwargs.update({
                'poolclass': StaticPool,
                'connect_args': {'check_same_thread': False},
            })
            # Remove pool settings that don't apply to SQLite
            for key in ['pool_size', 'max_overflow', 'pool_timeout', 'pool_recycle']:
                engine_kwargs.pop(key, None)
        
        self._async_engine = create_async_engine(async_url, **engine_kwargs)
        
        self._async_session_maker = async_sessionmaker(
            bind=self._async_engine,
            class_=AsyncSession,
            autocommit=False,
            autoflush=False,
            expire_on_commit=False
        )
        
        logger.info(f"Initialized asynchronous database engine: {async_url}")
    
    @property
    def engine(self):
        """Get synchronous engine."""
        if self._engine is None:
            raise DatabaseError("Synchronous engine not initialized")
        return self._engine
    
    @property
    def async_engine(self) -> AsyncEngine:
        """Get asynchronous engine."""
        if self._async_engine is None:
            raise DatabaseError("Asynchronous engine not initialized")
        return self._async_engine
    
    async def create_tables(self):
        """Create database tables."""
        try:
            if self.is_async:
                async with self._async_engine.begin() as conn:
                    await conn.run_sync(Base.metadata.create_all)
            else:
                Base.metadata.create_all(bind=self._engine)
            
            logger.info("Database tables created successfully")
        except Exception as e:
            raise DatabaseError(f"Failed to create database tables: {e}")
    
    async def drop_tables(self):
        """Drop database tables."""
        try:
            if self.is_async:
                async with self._async_engine.begin() as conn:
                    await conn.run_sync(Base.metadata.drop_all)
            else:
                Base.metadata.drop_all(bind=self._engine)
            
            logger.info("Database tables dropped successfully")
        except Exception as e:
            raise DatabaseError(f"Failed to drop database tables: {e}")
    
    @contextmanager
    def get_session(self) -> Generator[Session, None, None]:
        """
        Get synchronous database session.
        
        Yields:
            Database session
        """
        if self._session_maker is None:
            raise DatabaseError("Synchronous session maker not initialized")
        
        session = self._session_maker()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"Database session error: {e}")
            raise DatabaseError(f"Database operation failed: {e}")
        finally:
            session.close()
    
    @asynccontextmanager
    async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        Get asynchronous database session.
        
        Yields:
            Async database session
        """
        if self._async_session_maker is None:
            raise DatabaseError("Asynchronous session maker not initialized")
        
        async with self._async_session_maker() as session:
            try:
                yield session
                await session.commit()
            except Exception as e:
                await session.rollback()
                logger.error(f"Database session error: {e}")
                raise DatabaseError(f"Database operation failed: {e}")
    
    async def close(self):
        """Close database connections."""
        if self._async_engine:
            await self._async_engine.dispose()
            logger.info("Closed asynchronous database engine")
        
        if self._engine:
            self._engine.dispose()
            logger.info("Closed synchronous database engine")
    
    def check_connection(self) -> bool:
        """
        Check database connection.
        
        Returns:
            True if connection is working
        """
        try:
            if self.is_async:
                # For async, we'd need to use asyncio.run, but that's complex here
                # For now, just check engine existence
                return self._async_engine is not None
            else:
                with self._engine.connect() as conn:
                    conn.execute("SELECT 1")
                return True
        except Exception as e:
            logger.error(f"Database connection check failed: {e}")
            return False
    
    def get_database_info(self) -> dict:
        """
        Get database information.
        
        Returns:
            Dictionary with database info
        """
        return {
            'url': self.database_url,
            'is_async': self.is_async,
            'engine_type': 'async' if self.is_async else 'sync',
            'pool_size': self.config.database.pool_size,
            'echo': self.config.database.echo,
        }


# Global database manager instance
_db_manager: Optional[DatabaseManager] = None


def get_database_manager() -> DatabaseManager:
    """
    Get global database manager instance.
    
    Returns:
        DatabaseManager instance
    """
    global _db_manager
    if _db_manager is None:
        _db_manager = DatabaseManager()
    return _db_manager


def set_database_manager(manager: DatabaseManager):
    """
    Set global database manager instance.
    
    Args:
        manager: DatabaseManager instance
    """
    global _db_manager
    _db_manager = manager


@contextmanager
def get_session() -> Generator[Session, None, None]:
    """
    Get database session using global manager.
    
    Yields:
        Database session
    """
    manager = get_database_manager()
    with manager.get_session() as session:
        yield session


@asynccontextmanager
async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    """
    Get async database session using global manager.
    
    Yields:
        Async database session
    """
    manager = get_database_manager()
    async with manager.get_async_session() as session:
        yield session


async def init_database(drop_existing: bool = False):
    """
    Initialize database tables.
    
    Args:
        drop_existing: Whether to drop existing tables first
    """
    manager = get_database_manager()
    
    if drop_existing:
        await manager.drop_tables()
    
    await manager.create_tables()
    logger.info("Database initialized successfully")


async def close_database():
    """Close database connections."""
    global _db_manager
    if _db_manager:
        await _db_manager.close()
        _db_manager = None